#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py "Genus species"
#It is supposed to give you the parent phylum of the query term

import sys
import re
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

f = open(sys.argv[1], "rU")

lines = f.readlines()

num = len(lines)

rx = re.compile('(\w+ \w+).*') #only keep genus and species names as search query
#rx = re.compile('(\w+).*') #only keep genus name as search query

#m1 = re.compile('(.*\[.*\].*)\[(.*)\]')
#m2 = re.compile('(.*)\[(.*)\]')

n = 0

while n < num:
    curr_line = lines[n].split('\t')
    curr_prodigal_id = curr_line[0]
    curr_locus_tag = curr_line[1]
    curr_feat_type = curr_line[2]
    curr_start = int(curr_line[3])
    curr_stop = int(curr_line[4])
    curr_frame = curr_line[5]
    curr_blast_status = curr_line[6]
    curr_blast_evalue = curr_line[7]
    curr_tophit_id = curr_line[8]
    curr_tophit_sim = curr_line[9]
    curr_tophit_acc = curr_line[10]
    curr_tophit_gi = curr_line[11]
    curr_hit_def = curr_line[17]
    curr_hit_org = curr_line[18]
    curr_cluster_def = curr_line[31]
    curr_cluster_org = curr_line[32]

    if curr_blast_status == "Yes":
        searchterm = curr_hit_org
        if searchterm != "":
            pattern = rx.match(searchterm)
            st = pattern.group(1)
            handle = Entrez.esearch(db="Taxonomy", term=st)
            record = Entrez.read(handle)
            taxid = record["IdList"]
            #print curr_prodigal_id, "\t", taxid[0]
            if len(taxid) != 0:
                h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")
                rec = Entrez.parse(h)
                for r in rec:
                    lineage = r.get("Lineage")
                    l = lineage.split(';')
                    print curr_prodigal_id, "\t", l[2].lstrip()

    n += 1

f.close()
